Bio Data Science University of Zurich

Transformers with PyTorch

Transformers with PyTorch

The goal of today’s practical is to train a mini-transformer to generate sentences and answer questions about Shakespeare texts. The exercises are based on the original transformers paper: https://proceedings.neurips.cc/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf

%pip install datasets

import datasets
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

1. Tiny Shakespeare

Load the tiny_shakespeare dataset from Hugging Face and tokenize the characters https://huggingface.co/datasets/karpathy/tiny_shakespeare

import torch
from torch.utils.data import TensorDataset, DataLoader

# ---------------------------------------------------------------- load text
# The Hugging Face copy is the primary source. If it is unavailable (no
# network, or a datasets version mismatch) we fall back to the raw .txt file
# so that the rest of the practical still works.
try:
    from datasets import load_dataset
    ds = load_dataset('karpathy/tiny_shakespeare')
    train_text = ds['train'][0]['text']
    val_text = ds['validation'][0]['text']
    if len(train_text) < 100_000:
        raise ValueError('unexpected dataset structure')
    print('Loaded from Hugging Face.')
except Exception as e:
    print(f'Hugging Face load failed ({type(e).__name__}: {e}).')
    print('Falling back to the raw text file.')
    import urllib.request
    url = ('https://raw.githubusercontent.com/karpathy/char-rnn/'
           'master/data/tinyshakespeare/input.txt')
    with urllib.request.urlopen(url) as r:
        full_text = r.read().decode('utf-8')
    split = int(0.9 * len(full_text))
    train_text, val_text = full_text[:split], full_text[split:]

print(f'train: {len(train_text):,} characters, validation: {len(val_text):,} characters')
print(train_text[:250])

# ------------------------------------------------------------ character vocab
# The vocabulary is built from the train split only.
vocabulary = sorted(set(train_text))
vocab_size = len(vocabulary)
stoi = {ch: i for i, ch in enumerate(vocabulary)}
itos = {i: ch for i, ch in enumerate(vocabulary)}

unseen = set(val_text) - set(vocabulary)
if unseen:
    print(f'Warning: dropping {len(unseen)} character(s) unseen in train: {unseen}')
    val_text = ''.join(c for c in val_text if c in stoi)


def encode(s):
    """string -> list of integer token ids"""
    return [stoi[c] for c in s]


def decode(ids):
    """iterable of token ids -> string"""
    return ''.join(itos[int(i)] for i in ids)


print(f'vocab_size = {vocab_size}')
print('vocabulary:', ''.join(vocabulary).replace('\n', '\\n'))

# ------------------------------------------------- (cur_char, next_char) pairs
seq_len = 100
batch_size = 32


def make_sequences(text, seq_len):
    """Cut `text` into non-overlapping (cur_char, next_char) sequence pairs.

    next_char is cur_char shifted one position to the left, so position t of
    next_char holds the character the model must predict from position t of
    cur_char. Both come back with shape (n_sequences, seq_len).
    """
    data = torch.tensor(encode(text), dtype=torch.long)
    n_seq = (len(data) - 1) // seq_len          # -1 leaves room for the shift
    cur_char = data[:n_seq * seq_len].view(n_seq, seq_len)
    next_char = data[1:n_seq * seq_len + 1].view(n_seq, seq_len)
    return cur_char, next_char


train_x, train_y = make_sequences(train_text, seq_len)
val_x, val_y = make_sequences(val_text, seq_len)

train_loader = DataLoader(TensorDataset(train_x, train_y),
                          batch_size=batch_size, shuffle=True, drop_last=True)
val_loader = DataLoader(TensorDataset(val_x, val_y),
                        batch_size=batch_size, shuffle=False)

print(f'train sequences: {tuple(train_x.shape)}, validation: {tuple(val_x.shape)}')

xb, yb = next(iter(train_loader))
print(f'one batch -> cur_char {tuple(xb.shape)}, next_char {tuple(yb.shape)}')
print('cur_char [0][:60]:', repr(decode(xb[0][:60])))
print('next_char[0][:60]:', repr(decode(yb[0][:60])))

2. Positional encoding

Fill in the sinusoidal_positional_encoding function as described in the Attention Is All You Need paper

def sinusoidal_positional_encoding(max_len, d_model):
    # TODO: create a (max_len, d_model) matrix

    # TODO: generate positions 0,1,...,max_len-1

    # TODO: compute div_term (the exponential denominator)

    # TODO: assign sin to even positions and cos to odd positions

    pass

3. Scaled Dot-Product Attention

Fill in the dot-product function that is used to compute attention in a transformer

# TODO: Implement scaled dot-product attention
def scaled_dot_product_attention(Q, K, V, mask=None):
    # TODO: implement attention
    raise NotImplementedError("Implement the scaled dot-product attention")

4. Multi-Head Attention

Implement a Multi-Head attention block as described in the Attention Is All You Need paper

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=128, num_heads=8):
        super().__init__()
        assert d_model % num_heads == 0

        self.d_k = d_model // num_heads
        self.num_heads = num_heads

        # TODO: define linear layers W_q, W_k, W_v, W_o

    def forward(self, x):
        # TODO: implement multi-head attention forward pass
        raise NotImplementedError("Implement the MultiHeadAttention class")

5. Transformer Encoder Block

Define a transformer Encoder block that uses the attention layers you defined earlier

class TransformerEncoderBlock(nn.Module):
    def __init__(self, d_model=128, num_heads=8, dim_ff=512):
        super().__init__()
        # TODO: instantiate attention, layer norms, feed-forward block

    def forward(self, x):
        # TODO: implement encoder block logic
        raise NotImplementedError("Implement the TransformerEncoderBlock")

6. Mini Transformer

Build a transformer with all the parts you have defined in the previous exercises

class MiniTransformer(nn.Module):
    def __init__(self, vocab_size, d_model=128, n_layers=2):
        super().__init__()
        # TODO: define embeddings, positional encodings, encoder layers, final linear layer

    def forward(self, x):
        # TODO: implement forward pass
        raise NotImplementedError("Implement MiniTransformer")

7. Training

Train the transformer with the Shakespeare dataset that you loaded at the beginning

# TODO: write the training code

8. Testing the Transformer

Write the code that allows the transformer to give answers based on a string from the user